#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @File  : 11.使用PyTorch完成人脸关键点检测.py
# @Author: Richard Chiming Xu
# @Date  : 2021/11/17
# @Desc  :


import pandas as pd
import numpy as np
import torchvision.models
from PIL import Image, ImageDraw
from torch import nn
import torchvision.transforms as transforms
from torch.utils.data.dataset import Dataset
import torch
import torch.optim as optim
import time
from sklearn.metrics import mean_absolute_error

train_df = pd.read_csv('data/face/train.csv')
train_df = train_df.fillna(48)
train_img = np.load('data/face/train.npy')
test_img = np.load('data/face/test.npy')



torch.backends.cudnn.benchmark = False
# 判定是否使用GPU计算
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

config = {
    'BATCH_SIZE': 256,
    'LEARNING_RATE': 0.01,
    'EPOCHS': 50
}


class XunFeiDataset(Dataset):
    def __init__(self, img, keypoint, transform=None):
        self.img = img
        self.transform = transform
        self.keypoint = keypoint

    def __getitem__(self, index):
        img = Image.fromarray(self.img[:, :, index]).convert('RGB')

        if self.transform is not None:
            img = self.transform(img)

        return img, self.keypoint[index] / 96.0

    def __len__(self):
        return self.img.shape[-1]


train_loader = torch.utils.data.DataLoader(
    XunFeiDataset(train_img[:, :, :-500], train_df.values[:-500],
                    transforms.Compose([
                        transforms.ToTensor(),
        ])
    ),
    batch_size=10, shuffle=True
)

val_loader = torch.utils.data.DataLoader(
    XunFeiDataset(train_img[:, :, -500:], train_df.values[-500:],
                    transforms.Compose([
                        transforms.ToTensor(),
        ])
    ),
    batch_size=10, shuffle=False
)


# 搭建网络
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=5, padding=2),
            nn.BatchNorm2d(16),
            nn.ReLU())  # 16, 98, 98
        self.pool1 = nn.MaxPool2d(2)  # 16, 49, 49
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=3),
            nn.BatchNorm2d(32),
            nn.ReLU())  # 32, 46, 46
        self.layer3 = nn.Sequential(
            nn.Conv2d(32, 64, kernel_size=3),
            nn.BatchNorm2d(64),
            nn.ReLU())  # 64, 44, 44
        self.pool2 = nn.MaxPool2d(2)  # 64, 22, 22
        self.layer4 = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=3),
            nn.BatchNorm2d(128),
            nn.ReLU())  # 128, 20, 20
        self.pool3 = nn.MaxPool2d(2) # 128, 10, 10
        self.fc1 = nn.Linear(10 * 10 * 128, 128)
        self.fc2 = nn.Linear(128, 8)

    def forward(self, x):
        # x.shape = 10,3,96,96

        out = self.layer1(x) # out.shape = 10,16,96,96
        out = self.pool1(out) # out.shape = 10,16,48,48
        out = self.layer2(out) #out.shape = 10,32,46,46


        out = self.layer3(out) #out.shape = 10,64,44,44
        out = self.pool2(out) #out.shape = 10,64,22,22

        out = self.layer4(out) #out.shape = 10,128,20,20
        out = self.pool3(out) #out.shape = 10,128,10,10

        out = out.view(out.size(0), -1)
        out = self.fc1(out)
        out = self.fc2(out)
        return out


def train(train_loader, model, criterion, optimizer, epoch):
    model.train()

    for i, (input, target) in enumerate(train_loader):
        input = input.cuda(non_blocking=True).float()
        target = target.cuda(non_blocking=True).float()

        output = model(input)
        loss = criterion(output, target)
        optimizer.zero_grad()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        loss.backward()
        optimizer.step()

        if i % 200 == 0:
            print(loss.item())


def validate(val_loader, model):
    model.eval()

    val_feats = []
    with torch.no_grad():
        end = time.time()
        for i, (input, target) in enumerate(val_loader):
            input = input.cuda().float()
            target = target.cuda().float()
            output = model(input)
            val_feats.append(output.data.cpu().numpy())
    return val_feats


# CNN版本
# model = CNN()
# ResNet版本
# model = rn.ResNet(block=rn.BasicBlock, groups=[2, 2, 2, 2], num_classes=8)
model = torchvision.models.resnet18(pretrained=True)
model.fc = nn.Linear(model.fc.in_features,8,bias=False)

model.cuda(device)
criterion = nn.MSELoss().cuda()
optimizer = torch.optim.Adam(model.parameters(), 0.001)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.85)
best_acc = 0.0

for epoch in range(5):
    print('Epoch: ', epoch)

    train(train_loader, model, criterion, optimizer, epoch)

    val_feats = validate(val_loader, model)
    scheduler.step()

    val_feats = np.vstack(val_feats) * 96
    print('Val', mean_absolute_error(val_feats, train_df.values[-500:]))